import pymysql
import pymysql.cursors

class MySQLDB:
    def __init__(self, config):
        self.config = config
        self.connection = None
    
    def connect(self):
        """
        连接到MySQL数据库
        """
        try:
            self.connection = pymysql.connect(
                host=self.config['host'],
                port=self.config['port'],
                user=self.config['user'],
                password=self.config['password'],
                database=self.config['database'],
                cursorclass=pymysql.cursors.DictCursor,
                charset='utf8mb4'
            )
            return True
        except Exception as e:
            raise Exception(f"MySQL数据库连接失败: {e}")
    
    def disconnect(self):
        """
        断开数据库连接
        """
        if self.connection and self.connection.open:
            self.connection.close()
    
    def execute_query(self, query, params=None):
        """
        执行查询语句并返回结果
        """
        try:
            with self.connection.cursor() as cursor:
                cursor.execute(query, params)
                result = cursor.fetchall()
                return result
        except Exception as e:
            raise Exception(f"执行查询失败: {query}\n{params}\n{e}")
    
    def execute_update(self, query, params=None):
        """
        执行更新语句（INSERT, UPDATE, DELETE等）
        """
        try:
            with self.connection.cursor() as cursor:
                affected_rows = cursor.execute(query, params)
                self.connection.commit()
                return affected_rows
        except Exception as e:
            self.connection.rollback()
            raise Exception(f"执行更新失败: {query}\n{params}\n{e}")
    
    def get_table_structure(self, table_name):
        """
        获取表结构
        """
        query = """
        SELECT 
            COLUMN_NAME, 
            COLUMN_TYPE, 
            IS_NULLABLE, 
            COLUMN_KEY, 
            COLUMN_DEFAULT, 
            EXTRA, 
            COLUMN_COMMENT 
        FROM 
            INFORMATION_SCHEMA.COLUMNS 
        WHERE 
            TABLE_SCHEMA = %s 
            AND TABLE_NAME = %s
        ORDER BY 
            ORDINAL_POSITION
        """
        return self.execute_query(query, (self.config['database'], table_name))
    
    def get_all_tables(self):
        """
        获取所有表名
        """
        query = """
        SELECT 
            TABLE_NAME 
        FROM 
            INFORMATION_SCHEMA.TABLES 
        WHERE 
            TABLE_SCHEMA = %s 
            AND TABLE_TYPE = 'BASE TABLE'
        """
        result = self.execute_query(query, (self.config['database'],))
        return [row['TABLE_NAME'] for row in result]
    
    def get_all_views(self):
        """
        获取所有视图名
        """
        query = """
        SELECT 
            TABLE_NAME 
        FROM 
            INFORMATION_SCHEMA.TABLES 
        WHERE 
            TABLE_SCHEMA = %s 
            AND TABLE_TYPE = 'VIEW'
        """
        result = self.execute_query(query, (self.config['database'],))
        return [row['TABLE_NAME'] for row in result]
    
    def get_view_definition(self, view_name):
        """
        获取视图定义
        """
        # 注意：SHOW CREATE VIEW 语句不支持参数化查询，需要直接拼接视图名并用反引号括起来
        query = f"SHOW CREATE VIEW `{view_name}`"
        result = self.execute_query(query)
        if result and len(result) > 0:
            return result[0]['Create View']
        return None
        
    def get_table_comment(self, table_name):
        """
        获取表注释
        """
        query = """
        SELECT 
            TABLE_COMMENT 
        FROM 
            INFORMATION_SCHEMA.TABLES 
        WHERE 
            TABLE_SCHEMA = %s 
            AND TABLE_NAME = %s
        """
        result = self.execute_query(query, (self.config['database'], table_name))
        if result and len(result) > 0:
            return result[0]['TABLE_COMMENT']
        return None
    
    def get_table_count(self, table_name):
        """
        获取表的记录数
        """
        # 使用反引号括起来表名，以处理保留关键字
        query = f"SELECT COUNT(*) AS count FROM `{table_name}`"
        result = self.execute_query(query)
        if result and len(result) > 0:
            return result[0]['count']
        return 0
    
    def get_table_data_iterator(self, table_name, batch_size=1000, start_row=0):
        """
        获取表数据的迭代器，用于批量获取数据
        """
        try:
            with self.connection.cursor() as cursor:
                total_rows = self.get_table_count(table_name)
                current_row = start_row
                
                while current_row < total_rows:
                    # 使用反引号括起来表名，以处理保留关键字
                    query = f"SELECT * FROM `{table_name}` LIMIT %s OFFSET %s"
                    cursor.execute(query, (batch_size, current_row))
                    batch_data = cursor.fetchall()
                    
                    if not batch_data:
                        break
                    
                    yield batch_data
                    current_row += len(batch_data)
        except Exception as e:
            raise Exception(f"获取表数据失败: {table_name}\n{e}")
    
    def get_table_indexes(self, table_name):
        """
        获取表的所有索引信息
        """
        query = """
        SELECT 
            INDEX_NAME, 
            NON_UNIQUE, 
            COLUMN_NAME 
        FROM 
            INFORMATION_SCHEMA.STATISTICS 
        WHERE 
            TABLE_SCHEMA = %s 
            AND TABLE_NAME = %s
            AND INDEX_NAME != 'PRIMARY'  # 排除主键索引，主键已经单独处理
        ORDER BY 
            INDEX_NAME, SEQ_IN_INDEX
        """
        result = self.execute_query(query, (self.config['database'], table_name))
        
        # 按索引名分组，处理复合索引
        indexes = {}
        for row in result:
            index_name = row['INDEX_NAME']
            if index_name not in indexes:
                indexes[index_name] = {
                    'name': index_name,
                    'is_unique': not row['NON_UNIQUE'],
                    'columns': []
                }
            indexes[index_name]['columns'].append(row['COLUMN_NAME'])
        
        return list(indexes.values())